# 导入 numpy 库
import numpy as np
import cupy as cp
import random
import numba
import os
os.environ['NUMBA_WARNINGS'] = '0'

class KMeansPlusPlus:
    # 初始化类，传入聚类个数和数据
    def __init__(self, knum, data):
        self.k = knum # 聚类个数
        self.data = data # 数据
        self.n = len(data) # 数据个数
        self.centers = [] # 质心列表
        self.labels = np.zeros(self.n) # 标签列表
        self.sse = 0 # 平方误差和
        self.choose_centers()

    def choose_centers(self):
        # 随机选择第一个质心
        first = random.choice(self.data)
        self.centers.append(first)
        # 循环选择剩余的质心
        for i in range(1, self.k):
            # 计算每个数据点到最近质心的距离
            dist = cp.array([np.min([cp.linalg.norm(x - i).get() for i in self.centers]) for x in self.data])
            prob = cp.power(dist, 2) / cp.sum(dist)
            prob = prob / cp.sum(prob)
            prob = prob.get()
            next = np.random.choice(self.data, p=prob)
            self.centers.append(next)



    # 聚类
    def fit(self):
            # 计算每个数据点的标签
            for i in range(self.n):
                # 计算到每个质心的距离
                dist = [np.linalg.norm(self.data[i] - c) for c in self.centers]
                # 选择最近的质心作为标签 
                self.labels[i] = np.argmin(dist)
            # 更新质心
            new_centers = []
            for j in range(self.k):
                # 找出属于该类的数据点
                cluster = np.array(self.data)[self.labels == j]
                # 计算该类的均值作为新的质心
                new_center = np.mean(cluster, axis=0)
                new_centers.append(new_center)
            self.centers = new_centers

    # 计算平方误差和
    def getsse(self):
        # 初始化平方误差和为0
        self.sse = 0
        # 遍历每个数据点
        for i in range(self.n):
            # 计算到质心的距离的平方
            dist = np.linalg.norm(self.data[i] - self.centers[int(self.labels[i])])**2
            # 累加到平方误差和
            self.sse += dist
        # 返回平方误差和
        return self.sse

if __name__ == "__main__":
        data=np.array([1,5,4,6,9,102,45])
        a=KMeansPlusPlus(2,data)
        print(a.centers)
        a.fit()
        print(a.centers)
        a.fit()
        print(a.centers)
        #print(a.labels)
        print(a.getsse())
